from fastapi import WebSocket
import json
import asyncio
from typing import Dict, Any
from .ws_handler import WebSocketManager
from .mock_data import MockDataGenerator
from config.settings import DEFAULT_SYMBOL, DEFAULT_INST_TYPE, UPDATE_INTERVAL
from utils.logger import logger
from datetime import datetime
from cache.order_cache import order_cache
import sqlite3

ws_manager = WebSocketManager()
mock_generator = MockDataGenerator()

async def handle_public_ws(websocket: WebSocket):
    await ws_manager.connect(websocket)
    client_id = id(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            logger.info(f"[WebSocket接收] 客户端ID: {client_id}, 通道: public, 消息内容: {data}")
            
            # 处理字符串形式的ping
            if data == "ping":
                logger.info(f"[WebSocket处理] 客户端ID: {client_id}, 通道: public, 收到字符串PING")
                await ws_manager.handle_ping(websocket)
                continue
                
            # 处理JSON消息
            try:
                message = json.loads(data)
                
                # 处理 Bitget 格式的订阅消息 ("event": "subscribe")
                if message.get("event") == "subscribe":
                    await handle_subscription(websocket, message)
                
                # 兼容旧格式 ("op": "subscribe")
                elif message.get("op") == "subscribe":
                    logger.info(f"[WebSocket解析] 客户端ID: {client_id}, 通道: public, 消息类型: subscribe")
                    subscription = message.get("args", [])[0]
                    await handle_subscription(websocket, {"arg": subscription})
                elif message.get("op") == "ping":
                    logger.info(f"[WebSocket处理] 客户端ID: {client_id}, 通道: public, 收到JSON PING")
                    await ws_manager.handle_ping(websocket)
            except json.JSONDecodeError as e:
                logger.error(f"[WebSocket错误] 客户端ID: {client_id}, 通道: public, JSON解析错误: {str(e)}, 原始数据: {data}")
            except Exception as e:
                logger.error(f"[WebSocket错误] 客户端ID: {client_id}, 通道: public, 错误: {str(e)}")
                break
    except Exception as e:
        logger.error(f"[WebSocket连接错误] 客户端ID: {client_id}, 通道: public, 错误: {str(e)}")
    finally:
        ws_manager.disconnect(websocket)

async def handle_private_ws(websocket: WebSocket):
    await ws_manager.connect(websocket)
    client_id = id(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            logger.info(f"[WebSocket接收] 客户端ID: {client_id}, 通道: private, 消息内容: {data}")
            
            # 处理字符串形式的ping
            if data == "ping":
                logger.info(f"[WebSocket处理] 客户端ID: {client_id}, 通道: private, 收到字符串PING")
                await ws_manager.handle_ping(websocket)
                continue
            if data == "up":
                try:
                    logger.info(f"[WebSocket请求] 客户端ID: {client_id}, 收到刷新订单缓存请求")
                    # 执行刷新缓存操作
                    order_cache.refresh()
                    logger.info(f"[WebSocket响应] 客户端ID: {client_id}, 订单缓存已刷新")
                    
                    # 向客户端发送刷新成功的响应
                    response = {
                        "event": "refresh_cache",
                        "code": "0",
                        "msg": "刷新订单缓存成功",
                        "timestamp": int(datetime.now().timestamp() * 1000)
                    }
                    await websocket.send_json(response)
                except sqlite3.OperationalError as e:
                    # 特别处理SQLite线程错误
                    if "thread" in str(e).lower() or "locked" in str(e).lower():
                        logger.error(f"[WebSocket错误] 客户端ID: {client_id}, SQLite线程或锁定错误: {str(e)}")
                        error_response = {
                            "event": "refresh_cache",
                            "code": "1",
                            "msg": "数据库线程错误 - 请重试",
                            "timestamp": int(datetime.now().timestamp() * 1000)
                        }
                    else:
                        logger.error(f"[WebSocket错误] 客户端ID: {client_id}, 数据库操作错误: {str(e)}")
                        error_response = {
                            "event": "refresh_cache",
                            "code": "1",
                            "msg": f"数据库操作错误: {str(e)}",
                            "timestamp": int(datetime.now().timestamp() * 1000)
                        }
                    await websocket.send_json(error_response)
                except Exception as e:
                    if isinstance(e, ConnectionError):
                        logger.error(f"[WebSocket错误] 客户端ID: {client_id}, 数据库连接失败: {str(e)}")
                        error_response = {
                            "event": "refresh_cache",
                            "code": "1",
                            "msg": "数据库未连接 - 请先在GUI中连接数据库",
                            "timestamp": int(datetime.now().timestamp() * 1000)
                        }
                    else:
                        logger.error(f"[WebSocket错误] 客户端ID: {client_id}, 刷新订单缓存失败: {str(e)}")
                        # 向客户端发送刷新失败的响应
                        error_response = {
                            "event": "refresh_cache",
                            "code": "1",
                            "msg": f"刷新订单缓存失败: {str(e)}",
                            "timestamp": int(datetime.now().timestamp() * 1000)
                        }
                    await websocket.send_json(error_response)
                continue
            # 处理JSON消息
            try:
                message = json.loads(data)
                
                # 处理 Bitget 格式的订阅消息 ("event": "subscribe")
                if message.get("event") == "subscribe":
                    subscription = message.get("arg", {})
                    logger.info(f"[WebSocket订阅详情] Bitget格式订阅arg内容: {subscription}")
                    
                    channel = subscription.get("channel", "")
                    inst_id = subscription.get("instId", DEFAULT_SYMBOL)
                    inst_type = subscription.get("instType", DEFAULT_INST_TYPE)
                    
                    logger.info(f"[WebSocket订阅详情] Bitget格式提取到的channel: {channel}")
                    
                    # 发送订阅确认
                    response = {
                        "event": "subscribe",
                        "arg": subscription,
                        "code": "0",
                        "msg": "Subscribe success"
                    }
                    await websocket.send_json(response)
                    logger.info(f"[WebSocket处理] 客户端ID: {client_id}, 通道: private, 订阅确认: {channel}")
                    
                    # 无论ws_manager是否运行，都注册订阅者
                    if channel == "orders":
                        logger.info(f"[WebSocket订阅] Bitget格式发现orders频道订阅, 添加客户端ID: {client_id}到订阅列表")
                        ws_manager.add_order_subscriber(websocket)
                        logger.info(f"[WebSocket订阅] 当前orders频道订阅者数量: {len(ws_manager.order_subscribers)}")
                    
                    # 启动数据发送任务，仅在ws_manager运行时
                    if ws_manager.is_running and channel == "orders":
                        asyncio.create_task(send_mock_orders(websocket, subscription))
                elif message.get("op") == "login":
                    await ws_manager.handle_login(websocket)
                elif message.get("op") == "subscribe":
                    logger.info(f"[WebSocket解析] 客户端ID: {client_id}, 通道: private, 消息类型: subscribe")
                    
                    # 更加健壮的处理方式，适应不同格式的订阅消息
                    args = message.get("args", [])
                    logger.info(f"[WebSocket订阅详情] args内容: {args}")
                    
                    # 确保args是一个非空列表
                    if args and isinstance(args, list) and len(args) > 0:
                        subscription = args[0]
                        logger.info(f"[WebSocket订阅详情] subscription内容: {subscription}")
                        
                        channel = subscription.get("channel", "")
                        logger.info(f"[WebSocket订阅详情] 提取到的channel: {channel}")
                        
                        # 发送订阅确认响应
                        response = {
                            "event": "subscribe",
                            "arg": subscription,
                            "code": "0",
                            "msg": "Subscribe success"
                        }
                        await websocket.send_json(response)
                        logger.info(f"[WebSocket处理] 客户端ID: {client_id}, 通道: private, 订阅确认: {channel}")
                        
                        # 无论ws_manager是否运行，都注册订阅者，确保后续能收到推送
                        if channel == "orders":
                            logger.info(f"[WebSocket订阅] 发现orders频道订阅, 添加客户端ID: {client_id}到订阅列表")
                            ws_manager.add_order_subscriber(websocket)
                            logger.info(f"[WebSocket订阅] 当前orders频道订阅者数量: {len(ws_manager.order_subscribers)}")
                        
                        # 仅在ws_manager运行时才发送模拟数据
                        if ws_manager.is_running:
                            asyncio.create_task(send_mock_orders(websocket, subscription))
                    else:
                        logger.warning(f"[WebSocket订阅错误] 客户端ID: {client_id}, args格式不正确或为空: {args}")
                        # 尝试发送一个通用的订阅确认
                        response = {
                            "event": "subscribe",
                            "arg": {},
                            "code": "0",
                            "msg": "Subscribe success but args format invalid"
                        }
                        await websocket.send_json(response)
                elif message.get("op") == "ping":
                    logger.info(f"[WebSocket处理] 客户端ID: {client_id}, 通道: private, 收到JSON PING")
                    await ws_manager.handle_ping(websocket)
            except json.JSONDecodeError as e:
                logger.error(f"[WebSocket错误] 客户端ID: {client_id}, 通道: private, JSON解析错误: {str(e)}, 原始数据: {data}")
            except Exception as e:
                logger.error(f"[WebSocket错误] 客户端ID: {client_id}, 通道: private, 错误: {str(e)}")
                break
    except Exception as e:
        logger.error(f"[WebSocket连接错误] 客户端ID: {client_id}, 通道: private, 错误: {str(e)}")
    finally:
        ws_manager.disconnect(websocket)

async def send_mock_orders(websocket: WebSocket, subscription: Dict[str, Any]):
    inst_id = subscription.get("instId", DEFAULT_SYMBOL)
    inst_type = subscription.get("instType", DEFAULT_INST_TYPE)
    client_id = id(websocket)
    
    logger.info(f"[WebSocket订阅] 客户端ID: {client_id}, 开始发送订单数据, instId: {inst_id}, instType: {inst_type}")


async def handle_subscription(websocket: WebSocket, data: Dict[str, Any]):
    """处理订阅请求"""
    subscription = data.get("arg", {})
    channel = subscription.get("channel")
    client_id = id(websocket)
    
    if channel == "orders":
        ws_manager.add_order_subscriber(websocket)
        logger.info(f"客户端 {client_id} 订阅了 orders 频道")
    elif channel == "ticker":
        ws_manager.add_ticker_subscriber(websocket)
        logger.info(f"客户端 {client_id} 订阅了 ticker 频道")
    else:
        logger.warning(f"未知的订阅频道: {channel}")
        await websocket.send_json({
            "event": "error",
            "msg": f"不支持的频道: {channel}"
        })
        return
        
    # 发送订阅确认
    await websocket.send_json({
        "event": "subscribe",
        "arg": subscription,
        "code": "0",
        "msg": "Subscribe success"
    }) 